# 声明：本代码仅供学习和研究目的使用。使用者应遵守以下原则：  
# 1. 不得用于任何商业用途。  
# 2. 使用时应遵守目标平台的使用条款和robots.txt规则。  
# 3. 不得进行大规模爬取或对平台造成运营干扰。  
# 4. 应合理控制请求频率，避免给目标平台带来不必要的负担。   
# 5. 不得用于任何非法或不当的用途。
#   
# 详细许可条款请参阅项目根目录下的LICENSE文件。  
# 使用本代码即表示您同意遵守上述原则和LICENSE中的所有条款。  


# -*- coding: utf-8 -*-
# @Author  : relakkes@gmail.com
# @Time    : 2024/1/14 16:58
# @Desc    : 小红书存储实现类
import asyncio
import csv
import json
import os
import pathlib
from typing import Dict

import aiofiles

import config
from base.base_crawler import AbstractStore
from tools import utils, words
from var import crawler_type_var, media_crawler_db_var
from db import init_db, AsyncMysqlDB


def calculate_number_of_files(file_store_path: str) -> int:
    """计算数据保存文件的前部分排序数字，支持每次运行代码不写到同一个文件中
    Args:
        file_store_path;
    Returns:
        file nums
    """
    if not os.path.exists(file_store_path):
        return 1
    try:
        return max([int(file_name.split("_")[0])for file_name in os.listdir(file_store_path)])+1
    except ValueError:
        return 1


class XhsCsvStoreImplement(AbstractStore):
    csv_store_path: str = "data/xhs"
    file_count:int=calculate_number_of_files(csv_store_path)

    def make_save_file_name(self, store_type: str) -> str:
        """
        make save file name by store type
        Args:
            store_type: contents or comments

        Returns: eg: data/xhs/search_comments_20240114.csv ...

        """
        return f"{self.csv_store_path}/{self.file_count}_{crawler_type_var.get()}_{store_type}_{utils.get_current_date()}.csv"

    async def save_data_to_csv(self, save_item: Dict, store_type: str):
        """
        Below is a simple way to save it in CSV format.
        Args:
            save_item:  save content dict info
            store_type: Save type contains content and comments（contents | comments）

        Returns: no returns

        """
        pathlib.Path(self.csv_store_path).mkdir(parents=True, exist_ok=True)
        save_file_name = self.make_save_file_name(store_type=store_type)
        async with aiofiles.open(save_file_name, mode='a+', encoding="utf-8-sig", newline="") as f:
            f.fileno()
            writer = csv.writer(f)
            if await f.tell() == 0:
                await writer.writerow(save_item.keys())
            await writer.writerow(save_item.values())

    async def store_content(self, content_item: Dict):
        """
        Xiaohongshu content CSV storage implementation
        Args:
            content_item: note item dict

        Returns:

        """
        await self.save_data_to_csv(save_item=content_item, store_type="contents")

    async def store_comment(self, comment_item: Dict):
        """
        Xiaohongshu comment CSV storage implementation
        Args:
            comment_item: comment item dict

        Returns:

        """
        await self.save_data_to_csv(save_item=comment_item, store_type="comments")

    async def store_creator(self, creator: Dict):
        """
        Xiaohongshu content CSV storage implementation
        Args:
            creator: creator dict

        Returns:

        """
        await self.save_data_to_csv(save_item=creator, store_type="creator")

    async def add_new_relation(self, relation_item: Dict):
        """
        小红书平台不支持获取用户关系数据，此方法仅为实现抽象接口
        Args:
            relation_item: 关系数据字典（在小红书平台不使用）

        Returns:
            None
        """
        # 小红书不支持获取用户关系数据，直接返回
        return


class XhsDbStoreImplement(AbstractStore):
    async def _ensure_db_connection(self):
        """
        确保数据库连接已初始化
        """
        if not media_crawler_db_var.get():
            await init_db()

    async def store_content(self, content_item: Dict):
        """
        Xiaohongshu content DB storage implementation
        Args:
            content_item: content item dict

        Returns:

        """
        await self._ensure_db_connection()
        from .xhs_store_sql import (add_new_content,
                                    query_content_by_content_id,
                                    update_content_by_content_id)
        
        # 移除不在数据库表中的字段
        content_to_save = content_item.copy()
        fields_to_remove = ['xsec_token', 'xsec_source']
        for field in fields_to_remove:
            if field in content_to_save:
                content_to_save.pop(field)
                
        note_id = content_to_save.get("note_id")
        note_detail: Dict = await query_content_by_content_id(content_id=note_id)
        if not note_detail:
            content_to_save["add_ts"] = utils.get_current_timestamp()
            content_to_save["last_modify_ts"] = utils.get_current_timestamp()
            await add_new_content(content_to_save)
        else:
            content_to_save["last_modify_ts"] = utils.get_current_timestamp()
            await update_content_by_content_id(note_id, content_item=content_to_save)

    async def store_comment(self, comment_item: Dict):
        """
        Xiaohongshu content DB storage implementation
        Args:
            comment_item: comment item dict

        Returns:

        """
        await self._ensure_db_connection()
        from .xhs_store_sql import (add_new_comment,
                                    query_comment_by_comment_id,
                                    update_comment_by_comment_id)
                                    
        try:
            # 移除不在数据库表中的字段，并确保必要字段存在
            comment_to_save = comment_item.copy()
            fields_to_remove = ['xsec_token', 'xsec_source']
            for field in fields_to_remove:
                if field in comment_to_save:
                    comment_to_save.pop(field)
                    
            comment_id = comment_to_save.get("comment_id")
            if not comment_id:
                utils.logger.error(f"[XhsDbStoreImplement.store_comment] No comment_id in comment: {comment_item}")
                return
                
            comment_detail: Dict = await query_comment_by_comment_id(comment_id=comment_id)
            if not comment_detail:
                comment_to_save["add_ts"] = utils.get_current_timestamp()
                comment_to_save["last_modify_ts"] = utils.get_current_timestamp()
                # 确保必要字段有默认值
                comment_to_save.setdefault("sub_comment_count", 0)
                comment_to_save.setdefault("like_count", 0)
                comment_to_save.setdefault("pictures", "")
                comment_to_save.setdefault("parent_comment_id", 0)
                
                utils.logger.info(f"[XhsDbStoreImplement.store_comment] Adding new comment: {comment_to_save}")
                await add_new_comment(comment_to_save)
            else:
                comment_to_save["last_modify_ts"] = utils.get_current_timestamp()
                utils.logger.info(f"[XhsDbStoreImplement.store_comment] Updating comment: {comment_to_save}")
                await update_comment_by_comment_id(comment_id, comment_item=comment_to_save)
                
        except Exception as e:
            utils.logger.error(f"[XhsDbStoreImplement.store_comment] Error storing comment: {str(e)}")
            utils.logger.error(f"[XhsDbStoreImplement.store_comment] Comment data: {comment_item}")

    async def store_creator(self, creator: Dict):
        """
        Xiaohongshu content DB storage implementation
        Args:
            creator: creator dict

        Returns:

        """
        await self._ensure_db_connection()
        from .xhs_store_sql import (add_new_creator, query_creator_by_user_id,
                                    update_creator_by_user_id)
                                    
        # 移除不在数据库表中的字段，并确保必要字段存在
        creator_to_save = creator.copy()
        fields_to_remove = ['xsec_token', 'xsec_source']
        for field in fields_to_remove:
            if field in creator_to_save:
                creator_to_save.pop(field)
                
        user_id = creator_to_save.get("user_id")
        user_detail: Dict = await query_creator_by_user_id(user_id)
        if not user_detail:
            creator_to_save["add_ts"] = utils.get_current_timestamp()
            creator_to_save["last_modify_ts"] = utils.get_current_timestamp()
            # 确保必要字段有默认值
            creator_to_save.setdefault("gender", "")
            creator_to_save.setdefault("follows", "0")
            creator_to_save.setdefault("fans", "0")
            creator_to_save.setdefault("interaction", "0")
            await add_new_creator(creator_to_save)
        else:
            creator_to_save["last_modify_ts"] = utils.get_current_timestamp()
            await update_creator_by_user_id(user_id, creator_to_save)

    async def add_new_relation(self, relation_item: Dict):
        """
        小红书平台不支持获取用户关系数据，此方法仅为实现抽象接口
        Args:
            relation_item: 关系数据字典（在小红书平台不使用）

        Returns:
            None
        """
        # 小红书不支持获取用户关系数据，直接返回
        return


class XhsJsonStoreImplement(AbstractStore):
    json_store_path: str = "data/xhs/json"
    words_store_path: str = "data/xhs/words"
    lock = asyncio.Lock()
    file_count:int=calculate_number_of_files(json_store_path)
    WordCloud = words.AsyncWordCloudGenerator()

    def make_save_file_name(self, store_type: str) -> (str,str):
        """
        make save file name by store type
        Args:
            store_type: Save type contains content and comments（contents | comments）

        Returns:

        """

        return (
            f"{self.json_store_path}/{crawler_type_var.get()}_{store_type}_{utils.get_current_date()}.json",
            f"{self.words_store_path}/{crawler_type_var.get()}_{store_type}_{utils.get_current_date()}"
        )

    async def save_data_to_json(self, save_item: Dict, store_type: str):
        """
        Below is a simple way to save it in json format.
        Args:
            save_item: save content dict info
            store_type: Save type contains content and comments（contents | comments）

        Returns:

        """
        pathlib.Path(self.json_store_path).mkdir(parents=True, exist_ok=True)
        pathlib.Path(self.words_store_path).mkdir(parents=True, exist_ok=True)
        save_file_name,words_file_name_prefix = self.make_save_file_name(store_type=store_type)
        save_data = []

        async with self.lock:
            if os.path.exists(save_file_name):
                async with aiofiles.open(save_file_name, 'r', encoding='utf-8') as file:
                    save_data = json.loads(await file.read())

            save_data.append(save_item)
            async with aiofiles.open(save_file_name, 'w', encoding='utf-8') as file:
                await file.write(json.dumps(save_data, ensure_ascii=False, indent=4))

            if config.ENABLE_GET_COMMENTS and config.ENABLE_GET_WORDCLOUD:
                try:
                    await self.WordCloud.generate_word_frequency_and_cloud(save_data, words_file_name_prefix)
                except:
                    pass
    async def store_content(self, content_item: Dict):
        """
        content JSON storage implementation
        Args:
            content_item:

        Returns:

        """
        await self.save_data_to_json(content_item, "contents")

    async def store_comment(self, comment_item: Dict):
        """
        comment JSON storage implementatio
        Args:
            comment_item:

        Returns:

        """
        await self.save_data_to_json(comment_item, "comments")

    async def store_creator(self, creator: Dict):
        """
        Xiaohongshu content JSON storage implementation
        Args:
            creator: creator dict

        Returns:

        """
        await self.save_data_to_json(creator, "creator")

    async def add_new_relation(self, relation_item: Dict):
        """
        小红书平台不支持获取用户关系数据，此方法仅为实现抽象接口
        Args:
            relation_item: 关系数据字典（在小红书平台不使用）

        Returns:
            None
        """
        # 小红书不支持获取用户关系数据，直接返回
        return

async def save_creator(user_id: str, creator: Dict):
    """
    保存小红书创作者
    Args:
        user_id: 用户ID
        creator: 创作者信息字典

    Returns:

    """
    if not creator or not isinstance(creator, dict):
        utils.logger.warning(f"[store.xhs.save_creator] Invalid creator data for user_id {user_id}")
        return
        
    user_info = creator.get('basicInfo', {})
    if not user_info:
        utils.logger.warning(f"[store.xhs.save_creator] No basic info found for user_id {user_id}")
        return

    # 初始化默认值
    follows = 0
    fans = 0
    interaction = 0
    
    # 安全地获取互动信息
    interactions = creator.get('interactions', [])
    if interactions and isinstance(interactions, list):
        for i in interactions:
            if not isinstance(i, dict):
                continue
            if i.get('type') == 'follows':
                follows = i.get('count', 0)
            elif i.get('type') == 'fans':
                fans = i.get('count', 0)
            elif i.get('type') == 'interaction':
                interaction = i.get('count', 0)

    def get_gender(gender):
        if gender == 1:
            return '女'
        elif gender == 0:
            return '男'
        else:
            return ''

    try:
        local_db_item = {
            'user_id': user_id,  # 用户id
            'nickname': user_info.get('nickname', ''),  # 昵称
            'gender': get_gender(user_info.get('gender')), # 性别
            'avatar': user_info.get('images', ''), # 头像
            'desc': user_info.get('desc', ''), # 个人描述
            'ip_location': user_info.get('ipLocation', ''), # ip地址
            'follows': str(follows), # 关注数
            'fans': str(fans),  # 粉丝数
            'interaction': str(interaction), # 互动数
            'tag_list': json.dumps({tag.get('tagType', ''): tag.get('name', '') for tag in creator.get('tags', []) if isinstance(tag, dict)},
                                ensure_ascii=False), # 标签
            "last_modify_ts": utils.get_current_timestamp(), # 最后更新时间戳
        }
        utils.logger.info(f"[store.xhs.save_creator] creator:{local_db_item}")
        await XhsStoreFactory.create_store().store_creator(local_db_item)
        
    except Exception as e:
        utils.logger.error(f"[store.xhs.save_creator] Error saving creator {user_id}: {str(e)}")
        utils.logger.error(f"[store.xhs.save_creator] Creator data: {creator}")
